from common import np, pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import classification_report
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import IsolationForest
from sklearn.neighbors import LocalOutlierFactor
from sklearn.svm import OneClassSVM


def train_health_score_model(X, y):
    """训练健康度评分模型"""
    # 创建健康度评分 (0-100)
    # 正常状态为100分，故障状态根据故障类型扣分
    health_score = np.where(y == 0, 100, np.where(y == 1, 70, np.where(y == 2, 50, 30)))
    # 添加随机波动
    health_score += np.random.normal(0, 5, len(health_score))
    health_score = np.clip(health_score, 0, 100)

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, health_score, test_size=0.2, random_state=42)

    # 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 训练模型
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train_scaled, y_train)

    # 评估
    y_pred = model.predict(X_test_scaled)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    print(f"健康度评分模型RMSE: {rmse:.2f}")

    return model, scaler


def train_rul_model(X, y):
    """训练剩余使用寿命模型"""
    # 创建RUL标签（基于故障日志）
    # 假设设备设计寿命为1000小时
    rul = np.zeros(len(y))
    for i in range(len(y)):
        if y[i] == 0:
            # 正常状态，RUL递减
            rul[i] = max(0, 1000 - i)
        else:
            # 故障状态，RUL为0
            rul[i] = 0

    # 转换为时间窗口数据
    window_size = 10
    X_window, y_window = [], []
    for i in range(window_size, len(X)):
        X_window.append(X.iloc[i - window_size:i].values)
        y_window.append(rul[i])

    X_window = np.array(X_window)
    y_window = np.array(y_window)

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X_window, y_window, test_size=0.2, random_state=42)

    # 训练LSTM模型
    model = Sequential([
        LSTM(64, input_shape=(window_size, X_train.shape[2]), return_sequences=True),
        LSTM(32),
        Dense(16, activation='relu'),
        Dense(1)
    ])

    model.compile(optimizer='adam', loss='mse')
    model.fit(X_train, y_train, epochs=50, batch_size=32, validation_split=0.2)

    # 评估
    loss = model.evaluate(X_test, y_test)
    print(f"RUL模型测试损失: {loss:.2f}")

    return model


def train_fault_probability_model(X, y):
    """训练故障概率模型"""
    # 创建二分类标签（正常=0，故障=1）
    fault_binary = np.where(y > 0, 1, 0)

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, fault_binary, test_size=0.2, random_state=42)

    # 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 训练模型
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train_scaled, y_train)

    # 评估
    y_prob = model.predict_proba(X_test_scaled)[:, 1]
    auc = roc_auc_score(y_test, y_prob)
    print(f"故障概率模型AUC: {auc:.4f}")

    return model, scaler


def train_fault_type_model(X, y):
    """训练故障类型预测模型"""
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 训练模型
    model = GradientBoostingClassifier(n_estimators=100, random_state=42)
    model.fit(X_train_scaled, y_train)

    # 评估
    y_pred = model.predict(X_test_scaled)
    print("故障类型分类报告:")
    print(classification_report(y_test, y_pred))

    return model, scaler


def train_health_level_model(X, y):
    """训练健康等级分类模型"""
    # 创建健康等级标签
    # 0: 健康, 1: 轻微退化, 2: 中度退化, 3: 严重退化
    health_level = np.zeros(len(y))
    for i in range(len(y)):
        if y[i] == 0:
            # 随机添加退化状态
            if np.random.random() < 0.2:
                health_level[i] = np.random.choice([1, 2, 3], p=[0.6, 0.3, 0.1])
            else:
                health_level[i] = 0
        else:
            health_level[i] = 3  # 故障状态视为严重退化

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, health_level, test_size=0.2, random_state=42)

    # 标准化
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 训练模型
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train_scaled, y_train)

    # 评估
    accuracy = model.score(X_test_scaled, y_test)
    print(f"健康等级模型准确率: {accuracy:.4f}")

    return model, scaler


def train_anomaly_detection_model(X):
    """训练异常检测模型"""
    # 使用无监督方法
    models = {
        'isolation_forest': IsolationForest(contamination=0.05),
        'local_outlier': LocalOutlierFactor(novelty=True, contamination=0.05),
        'one_class_svm': OneClassSVM(nu=0.05)
    }

    # 标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 训练模型
    for name, model in models.items():
        if name != 'local_outlier':  # LOF需要特殊处理
            model.fit(X_scaled)

    return models, scaler


class DeviceHealthPredictor:
    def __init__(self):
        self.models = {}
        self.scalers = {}

    def load_models(self):
        """加载所有训练好的模型"""
        # 在实际应用中，这里会从文件加载预训练模型
        pass

    def predict_all(self, sensor_data):
        """综合预测所有健康指标"""
        # 预处理数据
        processed_data = self.preprocess_input(sensor_data)

        # 执行所有预测
        results = {
            'health_score': self.predict_health_score(processed_data),
            'rul': self.predict_rul(processed_data),
            'fault_probability': self.predict_fault_probability(processed_data),
            'fault_type': self.predict_fault_type(processed_data),
            'health_level': self.predict_health_level(processed_data),
            'anomaly_score': self.detect_anomaly(processed_data)
        }

        return results

    def preprocess_input(self, data):
        """预处理输入数据"""
        # 这里实现与训练时相同的预处理流程
        return processed_data

    def predict_health_score(self, data):
        """预测健康度评分"""
        scaled_data = self.scalers['health_score'].transform(data)
        return self.models['health_score'].predict(scaled_data)[0]

    def predict_rul(self, data):
        """预测剩余使用寿命"""
        # 转换为时间窗口格式
        window_data = self.create_window(data)
        return self.models['rul'].predict(window_data)[0][0]

    def predict_fault_probability(self, data):
        """预测故障概率"""
        scaled_data = self.scalers['fault_probability'].transform(data)
        return self.models['fault_probability'].predict_proba(scaled_data)[0][1]

    def predict_fault_type(self, data):
        """预测故障类型"""
        scaled_data = self.scalers['fault_type'].transform(data)
        return self.models['fault_type'].predict(scaled_data)[0]

    def predict_health_level(self, data):
        """预测健康等级"""
        scaled_data = self.scalers['health_level'].transform(data)
        return self.models['health_level'].predict(scaled_data)[0]

    def detect_anomaly(self, data):
        """检测异常"""
        scaled_data = self.scalers['anomaly'].transform(data)
        scores = {}
        for name, model in self.models['anomaly'].items():
            if name == 'local_outlier':
                scores[name] = model.score_samples(scaled_data.reshape(1, -1))[0]
            else:
                scores[name] = model.score_samples(scaled_data.reshape(1, -1))[0]
        return np.mean(list(scores.values()))


def display_health_dashboard(predictions):
    """显示健康度仪表盘"""
    print(f"设备健康度: {predictions['health_score']:.1f}/100")
    print(f"健康等级: {['健康', '轻微退化', '中度退化', '严重退化'][int(predictions['health_level'])]}")
    print(f"故障概率: {predictions['fault_probability'] * 100:.1f}%")

    if predictions['fault_type'] > 0:
        print(f"预测故障类型: 故障类型{predictions['fault_type']}")
    else:
        print("设备状态正常")

    print(f"剩余使用寿命: {predictions['rul']:.1f}小时")
    print(f"异常检测分数: {predictions['anomaly_score']:.2f} (越低越异常)")


def generate_maintenance_plan(predictions):
    """生成维护计划建议"""
    if predictions['health_score'] > 80:
        return "设备状态良好，按计划维护"
    elif predictions['health_score'] > 60:
        return f"建议在{predictions['rul'] / 24:.1f}天后进行预防性检查"
    elif predictions['health_score'] > 40:
        return "建议本周内安排维护"
    elif predictions['health_score'] > 20:
        return "建议48小时内安排维护"
    else:
        return "立即停机检修！"


def risk_alert_system(predictions):
    """风险预警系统"""
    alerts = []

    if predictions['fault_probability'] > 0.7:
        alerts.append(f"高风险故障预警 (概率: {predictions['fault_probability'] * 100:.1f}%)")

    if predictions['health_level'] == 3:
        alerts.append("严重退化状态，需立即关注")

    if predictions['anomaly_score'] < -0.5:
        alerts.append("检测到异常行为")

    if predictions['rul'] < 24:
        alerts.append(f"剩余寿命不足24小时")

    return alerts if alerts else ["状态正常"]


# 加载数据
data = load_data('sensor_data.csv', 'fault_log.csv')

# 预处理
X, y = preprocess_data(data)

# 特征工程
X = create_features(X)

# 训练所有模型
health_score_model, health_score_scaler = train_health_score_model(X, y)
rul_model = train_rul_model(X, y)
fault_prob_model, fault_prob_scaler = train_fault_probability_model(X, y)
fault_type_model, fault_type_scaler = train_fault_type_model(X, y)
health_level_model, health_level_scaler = train_health_level_model(X, y)
anomaly_models, anomaly_scaler = train_anomaly_detection_model(X)

# 保存模型
save_models({
    'health_score': (health_score_model, health_score_scaler),
    'rul': rul_model,
    'fault_probability': (fault_prob_model, fault_prob_scaler),
    'fault_type': (fault_type_model, fault_type_scaler),
    'health_level': (health_level_model, health_level_scaler),
    'anomaly': (anomaly_models, anomaly_scaler)
})

# 初始化预测器
predictor = DeviceHealthPredictor()
predictor.load_models()

# 实时预测循环
while True:
    # 获取最新传感器数据
    new_data = get_latest_sensor_data()

    # 预测
    predictions = predictor.predict_all(new_data)

    # 显示结果
    display_health_dashboard(predictions)

    # 生成维护建议
    maintenance_plan = generate_maintenance_plan(predictions)
    print(f"维护建议: {maintenance_plan}")

    # 检查风险
    alerts = risk_alert_system(predictions)
    for alert in alerts:
        print(f"⚠️ 警报: {alert}")

    # 等待下一分钟
    time.sleep(60)



    sensors = ['sensor_a', 'sensor_b', 'sensor_c', 'sensor_d']
    device_data = generate_data(sample_number, 0.3)

    device_data_copy = device_data.copy()
    sensors_copy = sensors.copy()

    for sensor in sensors:
        for window in [5, 10, 15]:
            device_data_copy[f'{sensor}_mean_{window}'] = device_data[sensor].rolling(window=window).mean()
            device_data_copy[f'{sensor}_std_{window}'] = device_data[sensor].rolling(window=window).std()
            device_data_copy[f'{sensor}_max_{window}'] = device_data[sensor].rolling(window=window).max()
            device_data_copy[f'{sensor}_min_{window}'] = device_data[sensor].rolling(window=window).min()

            sensors_copy.extend([
                f'{sensor}_mean_{window}',
                f'{sensor}_std_{window}',
                f'{sensor}_max_{window}',
                f'{sensor}_min_{window}'
            ])

    # 3. 变化特征
    for sensor in sensors:
        device_data_copy[f'{sensor}_diff_1'] = device_data[sensor].diff()
        device_data_copy[f'{sensor}_diff_3'] = device_data[sensor].diff(3)
        device_data_copy[f'{sensor}_pct_change'] = device_data[sensor].pct_change()
        sensors_copy.extend([f'{sensor}_diff_1', f'{sensor}_diff_3', f'{sensor}_pct_change'])

    # 4. 交互特征
    device_data_copy['sensor_ab_ratio'] = device_data['sensor_a'] / (device_data['sensor_b'] + 1e-8)
    device_data_copy['sensor_cd_sum'] = device_data['sensor_c'] + device_data['sensor_d']
    device_data_copy['sensor_product_ac'] = device_data['sensor_a'] * device_data['sensor_c']
    sensors_copy.extend(['sensor_ab_ratio', 'sensor_cd_sum', 'sensor_product_ac'])

    device_data_copy = device_data_copy.dropna()